import glob
import os

import cv2
import torch

from cjc.detector.detect_head import HeadDetector
from cjc.detector.detect_people import PeopleDetector
from cjc.detector.opt_tool import create_opt

detector = None


class ClearNonePersonPicture:

    def __init__(self, detector):
        self.detector = detector

    def detect_folder(self, folder):
        files = glob.glob(folder)
        print(files)
        for f in files:
            cvimg = cv2.imread(f)
            result = self.detector.detect(cvimg, f)
            if len(result) == 0:
                os.remove(f)
                print("+++++delete+++++", f)
            elif len(result) > 6:
                os.remove(f)
                print("+++++delete++>6+++", f)
            else:
                print("------pass-------", f)


if __name__ == '__main__':
    detector = PeopleDetector(create_opt())
    cleaner = ClearNonePersonPicture(detector)
    folder = '/mnt/d/datasets/face_data/smiles/*.jpg'
    cleaner.detect_folder(folder)
